from controller import Robot
import math
import numpy as np
import copy
import matplotlib.pyplot as plt

# create the Robot instance.
robot = Robot()

# get the time step of the current world.
timestep = int(robot.getBasicTimeStep())

# 机器人所有马达
#   0  1  2  3
#   LF RF RB LB
# 0 跨 ........
# 1 髋 ........
# 2 膝 ........
robot_motor = []  # LF RF RB LB
for i in range(4):
    robot_motor.append([])
# 左前足马达
robot_motor[0].append(robot.getDevice("LFJ0Motor"))
robot_motor[0].append(robot.getDevice("LFJ1Motor"))
robot_motor[0].append(robot.getDevice("LFJ2Motor"))
# 右前足马达
robot_motor[1].append(robot.getDevice("RFJ0Motor"))
robot_motor[1].append(robot.getDevice("RFJ1Motor"))
robot_motor[1].append(robot.getDevice("RFJ2Motor"))
# 右后足马达
robot_motor[2].append(robot.getDevice("RBJ0Motor"))
robot_motor[2].append(robot.getDevice("RBJ1Motor"))
robot_motor[2].append(robot.getDevice("RBJ2Motor"))
# 左后足马达
robot_motor[3].append(robot.getDevice("LBJ0Motor"))
robot_motor[3].append(robot.getDevice("LBJ1Motor"))
robot_motor[3].append(robot.getDevice("LBJ2Motor"))

# 机器人所有位置传感器
#   0  1  2  3
#   LF RF RB LB
# 0 跨 ........
# 1 髋 ........
# 2 膝 ........
robot_pos_sensor = []  # LF RF RB LB
for i in range(4):
    robot_pos_sensor.append([])
# 左前足位置传感器
robot_pos_sensor[0].append(robot.getDevice("LFJ0Sensor"))
robot_pos_sensor[0].append(robot.getDevice("LFJ1Sensor"))
robot_pos_sensor[0].append(robot.getDevice("LFJ2Sensor"))
# 右前足位置传感器
robot_pos_sensor[1].append(robot.getDevice("RFJ0Sensor"))
robot_pos_sensor[1].append(robot.getDevice("RFJ1Sensor"))
robot_pos_sensor[1].append(robot.getDevice("RFJ2Sensor"))
# 右后足位置传感器
robot_pos_sensor[2].append(robot.getDevice("RBJ0Sensor"))
robot_pos_sensor[2].append(robot.getDevice("RBJ1Sensor"))
robot_pos_sensor[2].append(robot.getDevice("RBJ2Sensor"))
# 左后足位置传感器
robot_pos_sensor[3].append(robot.getDevice("LBJ0Sensor"))
robot_pos_sensor[3].append(robot.getDevice("LBJ1Sensor"))
robot_pos_sensor[3].append(robot.getDevice("LBJ2Sensor"))

# 触地传感器
# LF RF RB LB
robot_touch_sensor = []
robot_touch_sensor.append(robot.getDevice("LF_Touch"))
robot_touch_sensor.append(robot.getDevice("RF_Touch"))
robot_touch_sensor.append(robot.getDevice("RB_Touch"))
robot_touch_sensor.append(robot.getDevice("LB_Touch"))

# 已测试
def webot_device_init():
    for leg in range(4):
        # 跨关节锁定
        robot_motor[leg][0].setPosition(0)
        # 初始化使能接触传感器
        robot_touch_sensor[leg].enable(timestep)
        # 初始化位置传感器
        for joint in range(3):
            robot_pos_sensor[leg][joint].enable(timestep)



# 顺序是 LF RF RB LB，单位是角度
def get_motor_angle(leg, name):
    angle = 0
    angle = robot_pos_sensor[leg][name].getValue()
    return angle * 180.0 / math.pi



def get_all_motor_angle():
    temp_list = list()
    for leg in range(4):  # 遍历四条腿 每条腿三个关节
        temp_list.append([])
        for joint in range(3):
            temp_list[leg].append(get_motor_angle(leg, joint))
    return temp_list


#设置电机角度
def set_motor_postion(leg, name, angle):
    robot_motor[leg][name].setPosition(angle)


def all_motor_stop():
    temp_list = get_all_motor_angle()
    for i in range(4):
        for j in range(3):
            set_motor_postion(i, j, temp_list[i][j] / 180 * math.pi)


class Quadruped_robot_mix():
    def __init__(self):
        self.zh = -0.2  # 抬腿到高度 0.4
        self.H_des = 0.3  # 0.5
        self.Tf = 0.5  # 飞行时间
        self.Ts = 0.5  # 支撑时间
        # 机器人机身参数
        self.l = 0.2 + 0.043 * 2  # 前后髋关节w
        self.w = 0.105  # 左右髋关节
        self.a0 = 0.0754  # 髋部连杆长度
        self.a1 = 0.2  # 大腿连杆长度
        self.a2 = 0.2  # 小腿连杆长度
        self.M = 20 + 8 * 1.0  # 躯干质量



    # 这里出来的坐标是世界系坐标，搞清楚电机转的顺序，这里都是顺时针为正转，腿的顺序，腿的顺序LF RF RB LB
    def forwardkinematics(self, theta):
        a0 = self.a0  # 腿长
        a1 = - self.a1  # 腿长
        a2 = - self.a2  # 腿长

        flag = [1.0, -1.0, -1.0, 1.0]
        # 这个和腿的顺序有关
        flag_d = [1.0, 1.0, -1.0, -1.0]
        # 这个和腿的顺序有关
        flag_l = [1.0, -1.0, -1.0, 1.0]
        # 足底位置

        temp_toe = list()
        for legIdx in range(len(theta)):
            s0 = math.sin(theta[legIdx][0] / 180.0 * math.pi)
            c0 = math.cos(theta[legIdx][0] / 180.0 * math.pi)
            s1 = math.sin(theta[legIdx][1] / 180.0 * math.pi)
            c1 = math.cos(theta[legIdx][1] / 180.0 * math.pi)

            s12 = math.sin(theta[legIdx][1] / 180.0 * math.pi + theta[legIdx][2] / 180.0 * math.pi)
            c12 = math.cos(theta[legIdx][1] / 180.0 * math.pi + theta[legIdx][2] / 180.0 * math.pi)

            temp_toe.append([])

            x = a1 * s1 + a2 * s12 + flag_d[legIdx] * self.l / 2
            y = flag[legIdx] * a0 * c0 - a1 * c1 * s0 - a2 * c12 * s0 + flag_l[legIdx] * self.w / 2
            z = flag[legIdx] * a0 * s0 + a1 * c1 * c0 + a2 * c12 * c0

            temp_toe[legIdx].append(x)
            temp_toe[legIdx].append(y)
            temp_toe[legIdx].append(z)

        print(temp_toe)
        return temp_toe

    def update_theta(self):
        self.theta = get_all_motor_angle()

    # 更新全部状态
    def update_robot_state(self):
        # 更新关节角
        self.update_theta()
        # 更新四条腿的运动学正解
        self.pos = self.forwardkinematics(self.theta)




    def run(self):
        # 初始化webots节点设备
        time = 0
        webot_device_init()
        while robot.step(timestep) != -1:
            # 更新所以数据
            time += 0.001
            self.update_robot_state()



if __name__ == '__main__':
    q = Quadruped_robot_mix()
    q.run()
